[Core] Refactor reference_counter out of memory store and plasma store #57590

Kunchd · 2025-10-09T04:02:21Z

Why are these changes needed?

As discovered in the PR to better define the interface for reference counter, plasma store provider and memory store both share thin dependencies on reference counter that can be refactored out. This will reduce entanglement in our code base and improve maintainability.

The main logic changes are located in

src/ray/core_worker/store_provider/plasma_store_provider.cc, where reference counter related logic is refactor into core worker
src/ray/core_worker/core_worker.cc, where factored out reference counter logic is resolved
src/ray/core_worker/store_provider/memory_store/memory_store.cc, where logic related to reference counter has either been removed due to the fact that it is tech debt or refactored into caller functions.

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run pre-commit jobs to lint the changes in this PR. (pre-commit setup)
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: davik <[email protected]>

Kunchd · 2025-10-09T17:55:27Z

src/ray/core_worker/store_provider/plasma_store_provider.cc

 }

 Status CoreWorkerPlasmaStoreProvider::Get(
-    const absl::flat_hash_set<ObjectID> &object_ids,


The logic below lifts the owner address resolution logic into the caller of Get.

There's a behavior change here. Before, we'd fetch the owner addresses for each batch. Now, we'll get them all up-front in the caller. It's probably fine. Keep an eye on the perf numbers to make sure there isn't a regression tho

Kunchd · 2025-10-09T17:56:15Z

src/ray/core_worker/core_worker.cc

  }
 }

+absl::flat_hash_map<ObjectID, rpc::Address> CoreWorker::GetObjectIdToOwnerAddressMap(


This function represents the owner address resolution logic lifted out of plasma store.

Kunchd · 2025-10-09T17:58:36Z

src/ray/core_worker/store_provider/memory_store/memory_store.cc

      for (auto &get_request : get_requests) {
        get_request->Set(object_id, object_entry);
-        // If ref counting is enabled, override the removal behaviour.
-        if (get_request->ShouldRemoveObjects() && ref_counter_ == nullptr) {


Remove after get flag is getting removed since it stems from the dark ages when reference counting wasn't a thing.

Kunchd · 2025-10-09T17:59:27Z

src/ray/core_worker/store_provider/memory_store/memory_store.cc

        const ray::RayObject &object, const ObjectID &object_id)> object_allocator)
    : io_context_(io_context),
-      ref_counter_(counter),
+      reference_counting_(reference_counting),


All reference counter null checks for whether reference counting is enabled is lifted into the caller.

Kunchd · 2025-10-09T18:00:22Z

src/ray/core_worker/store_provider/memory_store/memory_store.cc

-void CoreWorkerMemoryStore::Put(const RayObject &object, const ObjectID &object_id) {
+void CoreWorkerMemoryStore::Put(const RayObject &object,
+                                const ObjectID &object_id,
+                                const bool has_reference) {


Memory store only uses reference counter for checking if it has reference to an object or if reference counting is enabled. Both is lifted into caller.

Kunchd · 2025-10-09T22:01:37Z

src/mock/ray/core_worker/memory_store.h

  explicit DefaultCoreWorkerMemoryStoreWithThread(
      std::unique_ptr<InstrumentedIOContextWithThread> io_context)
-      : CoreWorkerMemoryStore(io_context->GetIoService()),
+      : CoreWorkerMemoryStore(io_context->GetIoService(), /*reference_counting=*/false),


Suggested change

: CoreWorkerMemoryStore(io_context->GetIoService(), /*reference_counting=*/false),

: CoreWorkerMemoryStore(io_context->GetIoService(), /*reference_counting=*/true),

There's no reason for the default core worker memory store to be created with reference counting disabled.

If it's not used in any local mode tests, then you should set it to true.

dayshah

haven't looked at this too hard yet, but at first glance it looks like having address in that vector will make wait/get perf on large lists even worse than it already is bc the impact of the vector set copying madness will be even worse

Do we need address to be in there??

Signed-off-by: davik <[email protected]>

…inimize_ref_count_dep

Kunchd · 2025-10-10T23:58:52Z

haven't looked at this too hard yet, but at first glance it looks like having address in that vector will make wait/get perf on large lists even worse than it already is bc the impact of the vector set copying madness will be even worse

Do we need address to be in there??

Nope, definitely doesn't need to be in the same data structure. I've moved it out for performance.

jjyao · 2025-10-21T23:57:38Z

Needs to rebase

israbbani · 2025-10-23T15:07:38Z

src/mock/ray/core_worker/memory_store.h

  explicit DefaultCoreWorkerMemoryStoreWithThread(
      std::unique_ptr<InstrumentedIOContextWithThread> io_context)
-      : CoreWorkerMemoryStore(io_context->GetIoService()),
+      : CoreWorkerMemoryStore(io_context->GetIoService(), /*reference_counting=*/false),


If it's not used in any local mode tests, then you should set it to true.

israbbani · 2025-10-23T15:11:45Z

src/ray/core_worker/store_provider/plasma_store_provider.cc


 Status CoreWorkerPlasmaStoreProvider::GetObjectsFromPlasmaStore(
-    absl::flat_hash_set<ObjectID> &remaining,
+    absl::flat_hash_map<ObjectID, int64_t> &remaining_object_id_to_idx,


Why is this needed?

This is a performance optimization. Since we now take in both the owner id and address in Get() instead of resolving owner address dynamically, we need a way to associate owner_id with its owner address quickly.

src/ray/core_worker/store_provider/plasma_store_provider.cc

israbbani · 2025-10-23T15:29:44Z

src/ray/core_worker/store_provider/plasma_store_provider.cc

 }

 Status CoreWorkerPlasmaStoreProvider::Get(
-    const absl::flat_hash_set<ObjectID> &object_ids,


There's a behavior change here. Before, we'd fetch the owner addresses for each batch. Now, we'll get them all up-front in the caller. It's probably fine. Keep an eye on the perf numbers to make sure there isn't a regression tho

israbbani · 2025-10-23T15:30:05Z

src/ray/core_worker/store_provider/plasma_store_provider.cc

-  std::vector<ObjectID> batch_ids;
-
-  int64_t num_total_objects = static_cast<int64_t>(object_ids.size());
+  absl::flat_hash_map<ObjectID, int64_t> remaining_object_id_to_idx;


What's the index for?

See https://github.com/ray-project/ray/pull/57590/files#r2456795459

src/ray/core_worker/store_provider/memory_store/memory_store.h

src/ray/core_worker/task_submission/BUILD.bazel

israbbani · 2025-10-23T15:55:49Z

src/ray/core_worker/core_worker.cc

+  std::vector<ObjectID> object_ids = {return_id};
+  auto owner_addresses = reference_counter_->GetOwnerAddresses(object_ids);


Same comment as above about overloading.

#57590 (comment)

israbbani · 2025-10-23T15:56:46Z

src/ray/core_worker/core_worker.cc

+    std::vector<ObjectID> object_ids =
+        std::vector<ObjectID>(plasma_object_ids.begin(), plasma_object_ids.end());
+    auto owner_addresses = reference_counter_->GetOwnerAddresses(object_ids);


Why not overload GetOwnerAddresses instead of creating a vector?

Currently, GetOwnerAddresses associate the output owner address vector with the input object_ids by index. If we take in a set instead, the ordering is now ill-defined. The current change should not cause any performance changes here as we simply moved the object_id vectorization out of plasma_store_provider and into core_worker.cc https://github.com/ray-project/ray/blob/master/src/ray/core_worker/store_provider/plasma_store_provider.cc#L262.

From offline discussion, we've concluded that the current way of storing owner address as a vector is poor cardinality as most object ids likely share the same owner address. However, the change to reduce cardinality will require significant logic changes across multiple files. GetOwnerAddress will be overloaded in a future PR that addresses the cardinality issue.

israbbani · 2025-10-23T15:57:21Z

src/ray/core_worker/BUILD.bazel

    hdrs = ["store_provider/plasma_store_provider.h"],
    deps = [
        ":common",
        ":core_worker_context",


…e by default Signed-off-by: davik <[email protected]>

cursor · 2025-10-24T19:51:32Z

src/ray/core_worker/core_worker.cc

+    // Prepare object ids vector and owner addresses vector
+    std::vector<ObjectID> object_ids =
+        std::vector<ObjectID>(plasma_object_ids.begin(), plasma_object_ids.end());
+    auto owner_addresses = reference_counter_->GetOwnerAddresses(object_ids);


Bug: Avoid Unnecessary Vector Copies

Creating a vector from a set and then calling GetOwnerAddresses() is inefficient. The PR discussion mentions "Why not overload GetOwnerAddresses instead of creating a vector?" This creates an unnecessary copy of potentially many ObjectIDs. The same issue appears at lines 3093-3094 and 3365-3367, but those involve smaller vectors so the performance impact is less severe.

See #57590 (comment)

davik and others added 2 commits October 9, 2025 03:52

Refactor reference_counter out of memory store and plasma store

3b1f487

Signed-off-by: davik <[email protected]>

Merge branch 'master' into minimize_ref_count_dep

0cfe8e3

Kunchd commented Oct 9, 2025

View reviewed changes

Kunchd marked this pull request as ready for review October 9, 2025 19:59

Kunchd requested review from a team, SongGuyang, kfstorm and raulchen as code owners October 9, 2025 19:59

Kunchd requested a review from israbbani October 9, 2025 19:59

Kunchd added the go add ONLY when ready to merge, run all tests label Oct 9, 2025

Kunchd commented Oct 9, 2025

View reviewed changes

ray-gardener bot added the core Issues that should be addressed in Ray Core label Oct 10, 2025

dayshah reviewed Oct 10, 2025

View reviewed changes

davik added 2 commits October 10, 2025 23:56

Optimize Get/Wait perfomance in plasma store

4ac6478

Signed-off-by: davik <[email protected]>

Merge branch 'minimize_ref_count_dep' of github.com:Kunchd/ray into m…

f01685f

…inimize_ref_count_dep

This comment was marked as outdated.

Sign in to view

Kunchd requested a review from dayshah October 14, 2025 17:56

Kunchd requested review from a team, pcmoritz and thomasdesr as code owners October 22, 2025 18:29

Kunchd requested review from a team, aslonnie, edoakes, jjyao and richardliaw as code owners October 22, 2025 18:29

Merge branch 'master' into minimize_ref_count_dep

7b55e41

Kunchd force-pushed the minimize_ref_count_dep branch from 9d17cce to 7b55e41 Compare October 22, 2025 18:41

Kunchd removed request for a team, aslonnie, edoakes, jjyao, pcmoritz, richardliaw and thomasdesr October 22, 2025 18:46

israbbani reviewed Oct 23, 2025

View reviewed changes

Address comments to set reference counting as enabled for memory stor…

a14329e

…e by default Signed-off-by: davik <[email protected]>

cursor bot reviewed Oct 24, 2025

View reviewed changes

Kunchd requested a review from israbbani October 24, 2025 19:58

Kunchd assigned Kunchd and israbbani Nov 7, 2025

	: CoreWorkerMemoryStore(io_context->GetIoService(), /reference_counting=/false),
	: CoreWorkerMemoryStore(io_context->GetIoService(), /reference_counting=/true),

		std::vector<ObjectID> object_ids = {return_id};
		auto owner_addresses = reference_counter_->GetOwnerAddresses(object_ids);

[Core] Refactor reference_counter out of memory store and plasma store #57590

Are you sure you want to change the base?

[Core] Refactor reference_counter out of memory store and plasma store #57590

Uh oh!

Conversation

Kunchd commented Oct 9, 2025

Why are these changes needed?

Related issue number

Checks

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dayshah left a comment

Choose a reason for hiding this comment

Uh oh!

Kunchd commented Oct 10, 2025

Uh oh!

This comment was marked as outdated.

Uh oh!

jjyao commented Oct 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cursor bot Oct 24, 2025

Choose a reason for hiding this comment

Bug: Avoid Unnecessary Vector Copies

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants